#!/bin/bash
set -euo pipefail

# 脚本名称：create_rootfs.sh

# 功能说明：
# 本脚本用于创建一个包含基本配置和工具的根文件系统，并生成一个虚拟磁盘镜像。

# 主要流程包括：
# 1. 安装必要的依赖工具（如 debootstrap、qemu-user-static 等）。
# 2. 根据用户输入的配置，创建根文件系统。
# 3. 配置网络、用户、主机名等系统设置。
# 4. 生成虚拟磁盘镜像，并将根文件系统写入其中。

# 注意事项：
# - 需要 root 权限运行部分命令，如 debootstrap、qemu-img。
# - 在执行操作前，建议备份重要数据，以防万一。

# == 默认配置 ==
ARCH="arm64"  # 可选: armhf, arm64 等架构
DISTRO="bionic"  # 可选: focal, jammy 等 Ubuntu 版本
MIRRORSITE="https://mirrors.huaweicloud.com/ubuntu-ports"
COMPONENTS="main,restricted,universe,multiverse"
EXTRA_PKGS="sudo,apt-utils,bash-completion,locales,net-tools,openssh-server,vim"
ROOTFS_DIR="$(pwd)/rootfs"
VMDK_IMG="ubuntu-${DISTRO}-${ARCH}.img"
VMDK_SIZE="4G"

# == 帮助信息 ==
usage() {
    echo "Usage: $0 [options]"
    echo
    echo "Options:"
    echo "  -a, --arch         Set target architecture (default: $ARCH)"
    echo "  -d, --distro       Set Ubuntu version (default: $DISTRO)"
    echo "  -m, --mirror       Set mirror site (default: $MIRRORSITE)"
    echo "  -s, --size         Set virtual disk size (default: $VMDK_SIZE)"
    echo "  -p, --packages     Set extra packages to install (comma-separated)"
    echo "  -h, --help         Show this help message and exit"
    echo
    exit 0
}

# == 解析用户输入 ==
while [[ $# -gt 0 ]]; do
    case "$1" in
        -a|--arch)
            ARCH="$2"
            shift 2
            ;;
        -d|--distro)
            DISTRO="$2"
            shift 2
            ;;
        -m|--mirror)
            MIRRORSITE="$2"
            shift 2
            ;;
        -s|--size)
            VMDK_SIZE="$2"
            shift 2
            ;;
        -p|--packages)
            EXTRA_PKGS="$2"
            shift 2
            ;;
        -h|--help)
            usage
            ;;
        *)
            echo "Unknown option: $1"
            usage
            ;;
    esac
done

# == 动态生成文件名 ==
VMDK_IMG="ubuntu-${DISTRO}-${ARCH}.img"

# == 检查所需工具 ==
install_dependencies() {
    local required=(debootstrap qemu-user-static qemu-system parted)
    local missing=()

    for cmd in "${required[@]}"; do
        if ! command -v "$cmd" &>/dev/null; then
            missing+=("$cmd")
        fi
    done

    if [ ${#missing[@]} -gt 0 ]; then
        echo "缺少依赖: ${missing[*]}，正在安装..."
        sudo apt-get update -qq
        sudo apt-get install -y -qq debootstrap qemu-user-static qemu-system parted >/dev/null
    fi
}

# == 创建虚拟磁盘并挂载 ==
create_virtual_disk() {
    if [ -f "$VMDK_IMG" ]; then
        echo "Removing existing image: $VMDK_IMG"
        rm "$VMDK_IMG"
    fi

    echo "Creating virtual disk: $VMDK_IMG"
    qemu-img create -f raw "$VMDK_IMG" "$VMDK_SIZE"
    sudo parted -s "$VMDK_IMG" mklabel msdos
    sudo parted -s "$VMDK_IMG" mkpart primary ext4 0% 100%
    LOOP_DEV=$(sudo losetup --show --find --partscan "$VMDK_IMG")
    echo y | sudo mkfs.ext4 -L rootfs "${LOOP_DEV}p1"

    mkdir -p "$ROOTFS_DIR"
    sudo mount "${LOOP_DEV}p1" "$ROOTFS_DIR"
}

# == 清理函数 ==
cleanup() {
    echo "Cleaning up..."
    if mountpoint -q "$ROOTFS_DIR"; then
        sudo umount "$ROOTFS_DIR"
    fi
    if [ -n "$LOOP_DEV" ]; then
        sudo losetup -d "$LOOP_DEV"
    fi
    echo "Done."
}
trap cleanup EXIT ERR

# == 使用 qemu-debootstrap 创建根文件系统 ==
create_rootfs() {
    echo "Creating root filesystem with qemu-debootstrap..."
    sudo qemu-debootstrap --arch="$ARCH" --components="$COMPONENTS" \
        --include="$EXTRA_PKGS" "$DISTRO" "$ROOTFS_DIR" "$MIRRORSITE"
}

# == 配置根文件系统功能模块 ==
configure_network() {
    echo "Configuring network..."
    cat <<EOF | sudo tee "$ROOTFS_DIR/etc/network/interfaces"
auto lo
iface lo inet loopback

auto eth0
iface eth0 inet dhcp
EOF
}

configure_locale() {
    echo "Configuring locale..."
    sudo chroot "$ROOTFS_DIR" locale-gen en_US.UTF-8
}

configure_users() {
    echo "Configuring users and passwords..."
    sudo chroot "$ROOTFS_DIR" bash -c "
        echo 'root:root' | chpasswd
        useradd -m -s /bin/bash ubuntu
        echo 'ubuntu:ubuntu' | chpasswd
        usermod -aG sudo ubuntu
    "
}

configure_hostname() {
    echo "Configuring hostname and hosts file..."
    echo "$DISTRO-$ARCH" | sudo tee "$ROOTFS_DIR/etc/hostname"
    cat <<EOF | sudo tee "$ROOTFS_DIR/etc/hosts"
127.0.0.1 localhost
127.0.1.1 $DISTRO-$ARCH
EOF
}

configure_packages() {
    echo "Updating and upgrading packages..."
    sudo chroot "$ROOTFS_DIR" apt-get update
    sudo chroot "$ROOTFS_DIR" apt-get -y upgrade
}

# == 主流程 ==
main() {
    install_dependencies
    create_virtual_disk
    create_rootfs

    configure_network
    configure_locale
    configure_users
    configure_hostname
    configure_packages

    echo "Switching root to $ROOTFS_DIR with architecture $ARCH..."
    sudo chroot "$ROOTFS_DIR" /bin/bash --rcfile <(cat ~/.bashrc; echo 'PS1="(chroot)> "')

    echo "Cleaning up the root filesystem..."
    sudo chroot "$ROOTFS_DIR" apt-get -y autoremove
    sudo chroot "$ROOTFS_DIR" apt-get -y clean
}

main "$@"
